import numpy as np

def beam_element_stiffness(E, I, L):#计算单个梁单元的刚度矩阵
    """
        计算梁单元的局部坐标系中的刚度矩阵
        Args:
            E: 材料的杨氏模量 (Pa)
            I: 惯性矩 (m^4)
            L: 梁单元的长度 (m)
        Returns:
            k_local: 梁单元的局部刚度矩阵
        """
    k = (E * I / L**3) * np.array([#接收三个参数：E(弹性模量)、I(惯性矩)、L(梁单元的长度)
        [12, 6*L, -12, 6*L],
        [6*L, 4*L**2, -6*L, 2*L**2],
        [-12, -6*L, 12, -6*L],
        [6*L, 2*L**2, -6*L, 4*L**2]
    ])
    return k



def assemble_global_stiffness(E, I, L, n_elements):#组装整体刚度矩阵，接收四个参数：E(弹性模量)、I(惯性矩)、L(梁的长度)、 n_elements梁单元数量
    n_nodes = n_elements + 1#计算节点数量，节点数量为单元数量+1
    K = np.zeros((2*n_nodes, 2*n_nodes))#创建一个大小为 2*n_nodes x 2*n_nodes 的零矩阵 K，用于存储整体刚度矩阵。每个节点有两个自由度（通常为垂直和水平方向的位移），因此需要两倍节点数的矩阵。
    for i in range(n_elements):#开始一个循环，遍历所有的梁单元。
        k = beam_element_stiffness(E, I, L)#调用 beam_element_stiffness 函数计算单个梁单元的局部刚度矩阵 k。
        dof = [2*i, 2*i+1, 2*i+2, 2*i+3]#定义列表dof(自由度)，包含当前单元两个端点的四个自由度的索引，对于每个单元，第一个节点的自由度索引是 2*i 和 2*i+1，第二个节点的自由度索引是 2*i+2 和 2*i+3。
        for ii in range(4):#嵌套循环遍历局部刚度矩阵 k 的四个元素。
            for jj in range(4):
                K[dof[ii], dof[jj]] += k[ii, jj]#将局部刚度矩阵 k 的元素累加到整体刚度矩阵 K 的对应位置上。这是通过索引 dof[ii] 和 dof[jj] 来实现的，确保只累加到与当前单元相关的自由度上。
    return K#函数返回组装好的整体刚度矩阵 K。


def apply_boundary_conditions(K, F, bc):#施加边界条件，接收三个参数：K（整体刚度矩阵）、F（力向量）、bc（边界条件的字典）
    for node, val in bc.items():#遍历边界条件字典 bc 中的每一项。node 是字典中的键，表示受边界条件约束的节点编号；val 是字典中的值，表示该节点的约束位移值。
        dof = [node]#为当前受边界条件约束的节点 node 创建一个自由度列表 dof。每个节点有两个自由度（例如，水平和垂直方向的位移），所以这里计算出两个自由度的索引。
        for d in dof:#遍历当前节点的所有自由度。
            K[d, :] = 0#将刚度矩阵 K 中与受约束自由度 d 相关的行和列的所有元素设置为 0。这表示在该自由度上没有力的作用，即该自由度被固定。
            K[:, d] = 0
            K[d, d] = 1#将刚度矩阵 K 的对角线上与受约束自由度 d 相关的元素设置为 1。这通常用于表示约束条件，即位移等于边界条件字典中指定的值。
            F[d] = val#将力向量 F 中与受约束自由度 d 相关的元素设置为边界条件字典 bc 中指定的值 val。这表示在该自由度上施加了一个等值且相反的力来平衡约束力。
    return K, F#函数返回修改后的整体刚度矩阵 K 和力向量 F。

def finite_element_analysis(E, I, L, n_elements, F, bc):#进行有限元分析，得到节点的位移和转角。函数接收六个参数：E（弹性模量）、I（惯性矩）、L（每个梁单元的长度）、n_elements（梁单元的数量）、F（力向量，包含作用在结构上的外力）、bc（边界条件的字典）
    K = assemble_global_stiffness(E, I, L, n_elements)#组装整体刚度矩阵 K，这个矩阵是通过对所有梁单元的局部刚度矩阵进行组装得到的。
    K, F = apply_boundary_conditions(K, F, bc)#调用 apply_boundary_conditions 函数来施加边界条件。这个函数修改了整体刚度矩阵 K 和力向量 F，以反映结构的固定边界和约束。
    displacements = np.linalg.solve(K, F)#使用linalg.solve 函数来求解线性方程组 K * displacements = F。这里的 displacements 是一个向量，包含了结构所有自由度上的位移和转角。np.linalg.solve 函数要求刚度矩阵 K 是非奇异的（即，可以求逆），这意味着结构的刚度矩阵应该是良好的，没有冗余或依赖关系。
    return displacements#函数返回位移向量 displacements，它包含了结构在所有自由度上的位移和转角。

# 参数定义
E = 210e9  # 杨氏模量 (Pa)
I = 1.0e-6  # 惯性矩 (m^4)
L = 1.0  # 每个单元的长度 (m)
n_elements = 4  # 单元数量
F = np.zeros(2*(n_elements+1))  # 力向量
F[2] = -100  # 在第二个节点施加向下的力 (N)

bc = {0: 0, 2*n_elements: 0, 2*n_elements+1: 0}  # 边界条件：第一个和最后一个节点固定

# 有限元分析
displacements = finite_element_analysis(E, I, L, n_elements, F, bc)
print("节点位移和转角：")
print(displacements)
